#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import fire

from pymongo import MongoClient


def GetPossibility(alt: str, geno: dict):
    total_individual_number = 0
    total_alt_number = 0

    allele = alt

    for item in geno.keys():
        total_individual_number += geno[item]
        count_alt = item.count(allele)
        total_alt_number += count_alt * geno[item]

    total_individual_number = total_individual_number * 2

    return (total_alt_number, total_individual_number)


def Tiantan_Possibilities():
    client = MongoClient("10.177.177.39", 26017)
    db = client["hanvcf"]
    db.authenticate("han", "han2019")
    collection_genotypes = db["genotypes_20190801"]
    collection_snvs = db["snvs_20190801"]

    print("rsid\tchrom\tpos\tref\talt\taaf")
    infile = open("/Users/genesis/Downloads/131_SNPs.txt", "r")
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        chrom = splitline[0].strip()
        rsid = splitline[1].strip()
        pos = int(splitline[2].strip())
        # result = collection_snvs.find_one({"rsid": rsid}, {"chrom": 1, "pos": 1, "rsid": 1})
        # if result == None:
        #     result2 = collection_snvs.find_one({"chrom":chrom,"pos":pos},{"chrom": 1, "pos": 1, "rsid": 1})
        #     if result2!=None:
        #         print(result2)
        alt = 0
        individual = 0
        ref_allele = ""
        alt_allele = ""
        for result in collection_genotypes.find({"rsid": rsid},
                                                {"chrom": 1, "pos": 1, "rsid": 1, "geno": 1, "ref": 1, "alt": 1}):
            this_alt, this_individual = GetPossibility(result["alt"], result["geno"])
            if ref_allele == "":
                ref_allele = result["ref"]
            if alt_allele == "":
                alt_allele = result["alt"]
            alt += this_alt
            individual += this_individual
        if individual != 0:
            # print(rsid + "\t" + chrom + "\t" + str(pos) + "\t" +ref_allele+"\t"+alt_allele+"\t"+ str(alt / individual))
            print("%s\t%s\t%s\t%s\t%s\t%s" % (rsid, chrom, str(pos), ref_allele, alt_allele, str(alt / individual)))
    infile.close()

    client.close()


if __name__ == '__main__':
    Tiantan_Possibilities()
    # fire.Fire()
